# build_lastcheck.py
# usage: results 和questions 拼接后，通过tmp_template2sql.json, 生成tmp_lastchecklist文件
# python ./src/tools/build_lastcheck.py --questions "./data/test.jsonl" --results "./result/results.jsonl"


import argparse
import json
from collections import defaultdict
import os

# 导入 clean_question 函数
from clean_question import clean_question


def load_jsonl(file_path):
    """读取 JSONL 文件，返回 id -> record 的字典"""
    data = {}
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            record = json.loads(line)
            data[record['id']] = record
    return data


def load_template2sql_by_template(template_file):
    """加载 tmp_template2sql.json，以 template 为 key 构建索引"""
    template2records = defaultdict(list)
    if not os.path.exists(template_file):
        print(f"Warning: {template_file} not found. Skipping existing template check.")
        return template2records
    with open(template_file, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            record = json.loads(line)
            tpl = record.get('template')
            if tpl:
                template2records[tpl].append(record)
    return template2records


def main():
    parser = argparse.ArgumentParser(description="Build last check list by question template.")
    parser.add_argument('--questions', required=True, help='Path to questions JSONL file')
    parser.add_argument('--results', required=True, help='Path to results JSONL file')
    args = parser.parse_args()

    # 加载问题和结果
    print("Loading questions...")
    questions = load_jsonl(args.questions)

    print("Loading results...")
    results = load_jsonl(args.results)

    # 加载模板SQL库，用于判断是否已存在
    print("Loading template2sql...")
    template2existing = load_template2sql_by_template('./tmp_template2sql.json')

    # 分组容器：template_str -> list of records
    grouped = defaultdict(list)

    # 顺序扫描 results 文件
    for result_id, result_record in results.items():
        if result_id not in questions:
            print(f"Warning: ID {result_id} in results but not found in questions. Skipping.")
            continue

        question_record = questions[result_id]
        raw_question = question_record['question']
        level = question_record['level']
        sql = result_record['sql']

        # 调用 clean_question 获取模板
        try:
            question_template = clean_question(raw_question)
        except Exception as e:
            print(f"Error cleaning question ID {result_id}: {e}")
            question_template = "Unknown Template"

        # 构建合并记录
        merged_record = {
            'id': result_id,
            'question': raw_question,
            'sql': sql,
            'level': level
        }

        # 添加到分组
        grouped[question_template].append(merged_record)

    # 写入输出文件
    output_path = './tmp_lastchecklist.txt'
    with open(output_path, 'w', encoding='utf-8') as out_f:
        for template, records in grouped.items():
            # 写入模板行
            out_f.write(f"{template}\n")

            # 判断该 template 是否在 tmp_template2sql.json 中存在（通过 template 匹配）
            has_existing = template in template2existing

            # 先写 * 记录（来自历史数据）
            if has_existing:
                for old_record in template2existing[template]:
                    id_ = old_record['id']
                    level_ = old_record['level']
                    sql_ = old_record['sql']
                    q_ = old_record['question']
                    out_f.write(f"   *{{\"id\": {id_}, \"level\": {level_}, \"flag\": 0, \"sql\": {json.dumps(sql_, ensure_ascii=False)}, \"question\": {json.dumps(q_, ensure_ascii=False)}}}\n")

            # 再写当前 results 中的记录
            for rec in records:
                out_f.write(f"    {{\"id\": {rec['id']}, \"level\": {rec['level']}, \"flag\": 0, \"sql\": {json.dumps(rec['sql'], ensure_ascii=False)}, \"question\": {json.dumps(rec['question'], ensure_ascii=False)}}}\n")

    print(f"Done. Output written to {output_path}")


if __name__ == '__main__':
    main()